# -*- coding: utf-8 -*-
# @Function: 数据导入器
# @Description: 负责将Excel文件中的问答对导入到数据库，包括文件读取、分词处理和数据库插入
# @Usage: 被all.py导入使用，执行数据导入流程
# @Dependencies: Z_config.py, db_manager.py, excel_reader.py, tokenizer.py

import os
import pandas as pd
from typing import List, Tuple
from D_db.core.db_manager import DBManager
from ..core.tokenizer import Tokenizer
from D_db.core.config import RESERVOIR_MAPPING
from D_db.io.excel_reader import ExcelReader

class QAImporter:
    def __init__(self):
        """初始化导入器"""
        self.db_manager = DBManager()
        self.excel_reader = ExcelReader()
        self.tokenizer = Tokenizer()

    def import_qa_from_excel(self, file_path: str, category_id: str) -> bool:
        """处理单个Excel文件

        Args:
            file_path (str): Excel文件路径
            category_id (str): 类别ID

        Returns:
            bool: 是否成功
        """
        try:
            print(f"\n处理文件: {os.path.basename(file_path)}")
            
            # 步骤1: 读取Excel文件
            df = self.excel_reader.read_file(file_path)
            print("步骤1: Excel文件读取完成")
            
            # 步骤2: 解析QA数据
            qa_pairs = self.excel_reader.parse_qa_data(df)
            if not qa_pairs:
                print(f"文件 {os.path.basename(file_path)} 中没有找到有效的QA数据")
                return False
            print(f"步骤2: 解析出 {len(qa_pairs)} 个问答对")
            
            # 步骤3: 对问题进行分词
            processed_qa_pairs = []
            for question, answer, context in qa_pairs:
                token_data = self.tokenizer.tokenize_question(question)
                processed_qa_pairs.append((question, answer, token_data, context))
            print("步骤3: 问题分词处理完成")
            
            # 步骤4: 插入数据库
            self.db_manager.insert_qa_data(processed_qa_pairs, category_id)
            print("步骤4: 数据库插入完成")
            
            print(f"文件 {os.path.basename(file_path)} 处理成功")
            return True
            
        except Exception as e:
            print(f"处理文件 {os.path.basename(file_path)} 失败: {e}")
            return False

    def process_folder(self, folder_path: str) -> Tuple[int, int]:
        """处理文件夹中的所有Excel文件

        Args:
            folder_path (str): 文件夹路径

        Returns:
            Tuple[int, int]: (成功处理的文件数, 总文件数)
        """
        try:
            # 获取所有Excel文件
            excel_files = [f for f in os.listdir(folder_path) if f.endswith('.xlsx')]
            total_files = len(excel_files)
            success_count = 0

            for file_name in excel_files:
                file_path = os.path.join(folder_path, file_name)
                reservoir_name = file_name.replace('.xlsx', '')

                if reservoir_name in RESERVOIR_MAPPING:
                    category_id = RESERVOIR_MAPPING[reservoir_name]
                    if self.import_qa_from_excel(file_path, category_id):
                        success_count += 1
                else:
                    print(f"跳过文件 {file_name} - 未找到匹配的category_id")
            
            print(f"\nExcel导入结果: {success_count}/{total_files} 个文件成功导入")
            return success_count, total_files
            
        except Exception as e:
            print(f"处理文件夹失败: {e}")
            return 0, 0
        finally:
            # 连接的关闭由主脚本控制
            pass

    def run(self, data_path: str):
        """运行导入流程

        Args:
            csv_file_path (str): CSV文件路径

        Returns:
            bool: 是否成功
        """
        try:
            # 连接数据库
            self.db_manager.connect()
            
            # 创建表
            self.db_manager.create_table()
            
            # 处理CSV文件
            result = self.process_csv_file(csv_file_path)
            
            return result
            
        except Exception as e:
            print(f"处理指定CSV文件失败: {e}")
            return False
        finally:
            # 关闭数据库连接
            self.db_manager.close()